#include "precond.hpp"
// #include "iter_solver.hpp"

#define L_IDX(i, j) ( (((i)*(  (i) - 1 )   ) >> 1) + (j)       )
#define U_IDX(i, j) ( (((i)*(2*N - (i) + 1)) >> 1) + (j) - (i) )

/// @brief 稠密LU分解
/// @tparam idx_t 索引类型
/// @tparam data_t 存储数据的浮点类型
/// @param a 稠密的A矩阵
/// @param l 稠密的下三角部分
/// @param u 稠密的上三角部分
/// @param num_rows 行数
/// @param num_cols 列数
template<typename idx_t, typename data_t>
void dense_LU_decomp(const data_t * a, data_t * l, data_t * u, const idx_t num_rows, const idx_t num_cols)
{
    assert(num_rows == num_cols);
    const idx_t N = num_rows;
    // assume: a为N*N的二维数组（一维存储），l为1+2+...+N-1=N*(N-1)/2的长度的数组，u为1+2+...+N=(N+1)*N/2的长度的数组

    // data_t (*A)[N] = (data_t (*)[N]) a;

    data_t rbuf[N];
    for (idx_t i = 0; i < N; i++) {
        // 拷入当前行的值
        for (idx_t j = 0; j < N; j++)
            rbuf[j] = a[i * N + j];// rbuf[j] = A[i][j];
        // 逐个消去左边的元素
        for (idx_t lj = 0; lj < i; lj++) {
            if (rbuf[lj] != 0.0) {
                rbuf[lj] /= u[U_IDX(lj, lj)];
                // 欲消该元会对该元右边元素产生影响
                for (idx_t rj = lj + 1; rj < N; rj++) {
                    if (u[U_IDX(lj, rj)] != 0.0)
                        rbuf[rj] -= rbuf[lj] * u[U_IDX(lj, rj)];
                }
            }
        }
        // 拷回到L和U中
        for (idx_t j = 0; j < i; j++)
            l[L_IDX(i, j)] = rbuf[j];
        for (idx_t j = i; j < N; j++)
            u[U_IDX(i, j)] = rbuf[j];
    }
}

/// @brief 稠密的前代
/// @tparam idx_t 索引类型
/// @tparam calc_t 执行计算时的浮点类型
/// @param l 稠密的下三角部分
/// @param rhs 右端项
/// @param x 解向量
/// @param num_rows 行数
/// @param num_cols 列数
template<typename idx_t, typename calc_t>
void dense_forward(const calc_t * l, const calc_t * rhs, calc_t * x, const idx_t num_rows, const idx_t num_cols)
{
    assert(num_rows == num_cols);
    const idx_t N = num_rows;
    // assume：l为1+2+...+N-1=N*(N-1)/2的长度的数组
    for (idx_t i = 0; i < N; i++) {
        calc_t res = rhs[i];
        for (idx_t j = 0; j < i; j++)
            res -= l[L_IDX(i, j)] * x[j];
        x[i] = res;
    }
}

/// @brief 稠密的回代
/// @tparam idx_t 索引类型
/// @tparam calc_t 执行计算时的浮点类型
/// @param u 稠密的上三角部分
/// @param rhs 右端项
/// @param x 解向量
/// @param num_rows 行数
/// @param num_cols 列数
template<typename idx_t, typename calc_t>
void dense_backward(const calc_t * u, const calc_t * rhs, calc_t * x, const idx_t num_rows, const idx_t num_cols)
{
    assert(num_rows == num_cols);
    const idx_t N = num_rows;
    // assume：u为1+2+...+N=(N+1)*N/2的长度的数组
    for (idx_t i = N - 1; i >= 0; i--) {
        calc_t diag = u[U_IDX(i, i)];
        calc_t res  = rhs[i];
        for (idx_t j = i + 1; j < N; j++)
            res -= u[U_IDX(i, j)] * x[j];
        x[i] = res / diag;
    }
}
#undef L_IDX
#undef U_IDX

/// @brief 稠密的LU分解
/// @tparam idx_t 索引类型
/// @tparam data_t 存储数据的浮点类型
/// @tparam calc_t 执行计算时的浮点类型
/// @tparam setup_t setup所用的浮点类型
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class DenseLU: public Precond<idx_t, calc_t, setup_t>
{
public:
    idx_t num_procs = 0;// 在当前层的进程数
    idx_t my_pid = -1;// 本进程的进程号
    idx_t glb_nrows = 0;// 当前层的全局行数
    MyParCSRMatrix<idx_t, setup_t, setup_t>* A_calc = nullptr;
    calc_t * u_data = nullptr, * l_data = nullptr;// 存储分解后的L和U数据
    calc_t * glb_buf0 = nullptr, * glb_buf1 = nullptr;// 存储全局向量的缓冲区
    int * recv_cnts = nullptr, * displs = nullptr;// 进行AllGather时所需的接收数目和偏移量
    // GMRESSolver<idx_t, calc_t, calc_t, calc_t> * gmres = nullptr;

    /// @brief 构造函数
    DenseLU() {}
    /// @brief 析构函数
    ~DenseLU() {
        if (glb_buf0) { delete [] glb_buf0; glb_buf0 = nullptr; }
        if (glb_buf1) { delete [] glb_buf1; glb_buf1 = nullptr; }
        if (u_data) { delete [] u_data; u_data = nullptr; }
        if (l_data) { delete [] l_data; l_data = nullptr; }
        if (recv_cnts) { delete [] recv_cnts; recv_cnts = nullptr; }
        if (displs   ) { delete [] displs   ; displs    = nullptr; }
        // if (gmres)  { delete gmres; gmres = nullptr; }
    }
    
    /// @brief 建立函数
    /// @param par_A 调用者传进来的矩阵
    void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & par_A)
    {
        A_calc = & par_A;
        MPI_Comm_rank(par_A.comm, & my_pid);
        MPI_Comm_size(par_A.comm, & num_procs);
        glb_nrows = par_A.glb_nrows;
        // 组装本进程内的矩阵
        const idx_t loc_nnz = par_A.diag->nnz + par_A.offd->nnz;
        struct ColID_Val {
            idx_t cid = -1;
            setup_t val = 0.0;
        };
        const idx_t loc_nrows = par_A.diag->nrows;
        idx_t * loc_rpt = new idx_t [loc_nrows + 1];// 本进程的局部行偏移
        loc_rpt[0] = 0;
        // 计算本进程的局部行偏移时需要将对角部分和非对角部分都算入
        for (idx_t i = 0; i < loc_nrows; i++) {
            loc_rpt[i + 1] =    par_A.diag->row_ptr[i + 1] - par_A.diag->row_ptr[i]
                            +   par_A.offd->row_ptr[i + 1] - par_A.offd->row_ptr[i];
        }
        // 全局行偏移
        idx_t * glb_rpt = new idx_t [glb_nrows + 1]; glb_rpt[0] = 0;
        constexpr MPI_Datatype mpi_idx_type = (sizeof(idx_t) == 4) ? MPI_INT : MPI_LONG_LONG;
        recv_cnts = new int [num_procs];
        for (int p = 0; p < num_procs; p++)
            recv_cnts[p] = par_A.rows_partition[p + 1] - par_A.rows_partition[p];
        // 通信获得每个进程负责范围内有多少个非零元
        MPI_Allgatherv(loc_rpt + 1, loc_nrows, mpi_idx_type, glb_rpt + 1, recv_cnts, par_A.rows_partition, mpi_idx_type, par_A.comm);
        // 从而可以整理得到全局行偏移
        for (idx_t i = 0; i < glb_nrows; i++) {
            glb_rpt[i + 1] += glb_rpt[i];
        }
        const idx_t glb_nnz = glb_rpt[glb_nrows];
        
        for (idx_t i = 0; i < loc_nrows; i++) {
            loc_rpt[i + 1] += loc_rpt[i];
        }
        // 记录要从其它进程处接收的数目和位置偏移量
        displs = new int [num_procs];
        for (int p = 0; p < num_procs; p++) {
            displs   [p] = glb_rpt[ par_A.rows_partition[p    ] ];
            assert(0 <= displs[p]); assert(displs[p] < glb_nnz);

            recv_cnts[p] = glb_rpt[ par_A.rows_partition[p + 1] ] - displs[p];
            // if (recv_cnts[p] < 0) printf("Error ! Proc %d recv_cnts[%d] = %d, displs[%d] = %d glb_rpt[ par_A.rows_partition[p + 1] ] %d\n",
            //     my_pid, p, recv_cnts[p], p, displs[p], glb_rpt[ par_A.rows_partition[p + 1] ]);
            assert(0 <= recv_cnts[p]);

            displs   [p] *= sizeof(ColID_Val);
            recv_cnts[p] *= sizeof(ColID_Val);
        }
        
        ColID_Val * loc_nonzeros = new ColID_Val [loc_nnz];// 本地的非零元
        ColID_Val * glb_nonzeros = new ColID_Val [glb_nnz];// 全局的非零元
        // 将本进程负责行范围内的非零元归入（包括对角和非对角部分）
        assert(loc_rpt[loc_nrows] == loc_nnz);
        for (idx_t i = 0; i < loc_nrows; i++) {
            const idx_t offset = loc_rpt[i];
            idx_t cnt = 0;
            for (idx_t p = par_A.diag->row_ptr[i]; p < par_A.diag->row_ptr[i + 1]; p++) {
                loc_nonzeros[offset + cnt].cid = par_A.diag->col_idx[p] + par_A.beg_row;
                loc_nonzeros[offset + cnt].val = par_A.diag->vals   [p];
                cnt ++;
            }
            for (idx_t p = par_A.offd->row_ptr[i]; p < par_A.offd->row_ptr[i + 1]; p++) {
                loc_nonzeros[offset + cnt].cid = par_A.col_map_offd[par_A.offd->col_idx[p]];
                loc_nonzeros[offset + cnt].val = par_A.offd->vals   [p];
                cnt ++;
            }
            assert(cnt == loc_rpt[i + 1] - offset);
        }

        // printf("Proc %d recv_cnts %d %d %d\n", my_pid, recv_cnts[0], recv_cnts[1], recv_cnts[2]);
        // printf("Proc %d displs    %d %d %d\n", my_pid,    displs[0],    displs[1],    displs[2]);

        // 集合通信获得全局的矩阵信息
        MPI_Allgatherv(loc_nonzeros, loc_nnz * sizeof(ColID_Val), MPI_BYTE, glb_nonzeros, recv_cnts, displs, MPI_BYTE, par_A.comm);

        l_data = new calc_t [glb_nrows * (glb_nrows - 1) / 2];// 下三角
        u_data = new calc_t [glb_nrows * (glb_nrows + 1) / 2];// 上三角
        calc_t * dense_A = new calc_t [glb_nrows * glb_nrows];// 稠密的待分解矩阵
        glb_buf0 = new calc_t [glb_nrows];
        glb_buf1 = new calc_t [glb_nrows];

        // 对待分解的稠密矩阵进行非零元的填充
        for (idx_t i = 0; i < glb_nrows; i++) {
            for (idx_t j = 0; j < glb_nrows; j++) {
                dense_A[i * glb_nrows + j] = 0.0;
            }
            for (idx_t p = glb_rpt[i]; p < glb_rpt[i + 1]; p++) {
                const idx_t j = glb_nonzeros[p].cid;
                dense_A[i * glb_nrows + j] = glb_nonzeros[p].val;
            }
        }

        // 执行LU分解
        dense_LU_decomp(dense_A, l_data, u_data, glb_nrows, glb_nrows);

        delete [] loc_rpt; loc_rpt = nullptr;
        delete [] glb_rpt; glb_rpt = nullptr;
        delete [] loc_nonzeros; loc_nonzeros = nullptr;
        delete [] glb_nonzeros; glb_nonzeros = nullptr;

        // 将偏移量和收取数目恢复为向量的数目
        for (int p = 0; p < num_procs; p++) {
            displs   [p] = par_A.rows_partition[p];
            recv_cnts[p] = par_A.rows_partition[p + 1] - displs[p];
            // displs   [p] *= sizeof(calc_t);
            // recv_cnts[p] *= sizeof(calc_t);
        }
    }
    /// @brief 求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @param zero_guess 是否采取零初值优化 
    void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool zero_guess=true)
    {
        static_assert(sizeof(calc_t) == 4 || sizeof(calc_t) == 8);
        constexpr MPI_Datatype mpi_val_type = (sizeof(calc_t) == 4) ? MPI_FLOAT : MPI_DOUBLE;
        if (zero_guess) {
            assert(b.end_row - b.beg_row == recv_cnts[my_pid]);
            // 从所有进程处收集向量
            MPI_Allgatherv(b.data, b.end_row - b.beg_row, mpi_val_type,
                glb_buf0, recv_cnts, displs, mpi_val_type, b.comm);

            dense_forward (l_data, glb_buf0, glb_buf1, glb_nrows, glb_nrows);// 前代
            dense_backward(u_data, glb_buf1, glb_buf0, glb_nrows, glb_nrows);// 回代

            for (idx_t i = x.beg_row; i < x.end_row; i++)
                x.data[i - x.beg_row] = glb_buf0[i];
        }
    }
};